class RSTTree:
    def __init__(self) -> None:
        self.leaf_nodes = []
        self.all_nodes = []
        self.root_node = None
        self.branching_count = 0
        pass
    # def mapping(self):
    #     self.relation_mapping(self.root_node)
    #     return

    # def relation_mapping(self, node):
    #     if node.nuclearity != 'Root' and node.relation != "span":
    #         mapped_relation = rstdt_mapping18.get(node.relation.lower())
    #         assert mapped_relation is not None
    #         node.relation = mapped_relation
        
    #     for child_node in node.child_nodes:
    #         self.relation_mapping(child_node)
    #     return

    def save(self, file_path):
        with open(file_path, mode='w', encoding='utf8') as outf:
            self.save_node(self.root_node, outf)
        return
    
    def save_node(self, node, outf):
        format_str = ""
        
        for idx in range(node.depth):
            format_str += '\t'

        if node.is_leaf():
            node_str = 'leaf'
            id_str = str(node.edu_start)
            text_str = "(text _!" + " ".join(node.text) + "_!)"
        else:
            node_str = 'span'
            id_str = str(node.edu_start) + " " + str(node.edu_end)
        
        node_line = format_str + "( " + node.nuclearity  + " (" + node_str + " " + id_str + ")"

        if node.nuclearity != "Root": 
            node_line += " (rel2par " + node.relation + ")"

        if node.is_leaf():
            node_line += " " + text_str + " )"

        outf.write(node_line + "\n") 

        for child_node in node.child_nodes: 
            self.save_node(child_node, outf)
        
        if not node.is_leaf():
            outf.write(format_str + ")" + "\n")

        return

class RSTNode:
    def __init__(self) -> None:
        self.text = None
        self.child_nodes = []
        self.parent_node = None

        self.edu_start = -1
        self.edu_end = -1

        self.nuclearity = None
        self.relation = None

        self.depth = -1
    
    def is_leaf(self):
        if self.text != None:
            return True
        else:
            return False